import os.path

from flask import request

from owl_ai.domain.rag_entity import KnowledgeBase
from owl_ai.service.rag.rag_service import RAGKBService, RAGEmbeddingService
from owl_common.base.model import AjaxResponse
from owl_common.config import OWLConfig
from owl_common.descriptor.serializer import JsonSerializer
from owl_common.descriptor.validator import BodyValidator
from owl_common.utils import FileUploadUtil
from ... import reg


@reg.api.route("/ai/rag/knowledgebase/create", methods=["POST"])
@JsonSerializer()
@BodyValidator()
def create_rag_knowledge_base(dto: KnowledgeBase):
    knowledge_base_id = RAGKBService.create_knowledge_base(dto)
    return AjaxResponse.from_success(data=knowledge_base_id)


@reg.api.route("/ai/rag/knowledgebase/file/upload", methods=["POST"])
@JsonSerializer()
def knowledgebase_file_upload():
    knowledge_base_id = request.form.get('knowledge_base_id')
    upload_files = request.files
    if upload_files:
        file = upload_files.get("file")
        # 文件保存
        resource_path = FileUploadUtil.upload(file, OWLConfig().upload_path)
        # 文件分块并向量化
        documents = RAGKBService.file_split(os.path.join(OWLConfig().upload_path, resource_path))
        embeddings = RAGEmbeddingService.document_embedding_save("test", documents)

        # 记录存入数据库
        print(embeddings)
        return AjaxResponse.from_success(data=embeddings)


@reg.api.route("/ai/rag/knowledgebase/search", methods=["POST"])
@JsonSerializer()
def knowledgebase_search():
    query = request.args.get("query")
    documents = RAGEmbeddingService.document_embedding_search(db_name="test", query=query)
    # 搜索
    return AjaxResponse.from_success(data=documents)


@JsonSerializer()
@BodyValidator()
@reg.api.route("/ai/rag/knowledgebase/find/byId", methods=["POST"])
def find_by_id(dto: KnowledgeBase):
    knowledge_base = RAGKBService.find_by_id(dto.id)
    return AjaxResponse.from_success(data=knowledge_base)
